package com.learning.spring.cloud.alibaba.gateway.confg;

import org.reactivestreams.Publisher;
import org.springframework.cloud.gateway.filter.ratelimit.KeyResolver;
import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter;
import org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter;
import org.springframework.cloud.gateway.route.RouteLocator;
import org.springframework.cloud.gateway.route.builder.RouteLocatorBuilder;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

/**
 * ClassName: RequestRateLimiterConfig
 * Description: 限流  RequestRateLimiter 相关的配置
 * Date: 2021/9/13 16:05 【需求编号】
 *
 * @author Sam Sho
 * @version V1.0.0
 */
@Configuration
public class RequestRateLimiterConfig {

    /**
     * ModifyResponseBodyGatewayFilterFactory
     *
     * @param builder
     * @return
     */
    @Bean
    public RouteLocator routes(RouteLocatorBuilder builder) {
        return builder.routes()
                .route("order", r -> r.path("/order/**")
                        .filters(f -> f.stripPrefix(1)
                                .requestRateLimiter().rateLimiter(RedisRateLimiter.class, RequestRateLimiterConfig::configConsumer).and()
                                .modifyResponseBody(String.class, String.class,
                                        RequestRateLimiterConfig::apply)).uri("lb://hc-service-order").filters())
                .build();
    }

    private static void configConsumer(RedisRateLimiter.Config config) {
        config.setReplenishRate(1);
        config.setBurstCapacity(30);
    }

    private static Publisher<String> apply(ServerWebExchange exchange, String s) {
        final ServerHttpResponse response = exchange.getResponse();
        final HttpStatus statusCode = response.getStatusCode();
        System.out.println(statusCode);
        if (statusCode == HttpStatus.TOO_MANY_REQUESTS) {
            return Mono.just("限流了~~~");
        }
        return Mono.just(s);
    }

    /**
     * 1. 在 redis中会产生两个Key：
     * - request_rate_limiter.{}.tokens
     * - request_rate_limiter.{}.timestamp
     * <p>
     * 2. request_rate_limiter.lua 逻辑
     * - local fill_time = capacity/rate 计算填充满令牌桶需要的时间
     * - local ttl = math.floor(fill_time*2) 计算过期时间
     * -
     *
     * @return
     */
    @Bean
    RateLimiter myRateLimiter() {
        return new RedisRateLimiter(10, 20);
    }

    @Bean
    KeyResolver userKeyResolver() {
        return exchange -> Mono.just(exchange.getRequest().getQueryParams().getFirst("user"));
    }
}
